import torch
import torch.nn as nn

# 定义一个简单的 RNN 模型
class SimpleRNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleRNN, self).__init__()
        self.rnn = nn.RNN(input_size, hidden_size, num_layers=4,batch_first=True)  # 定义 RNN 层
        self.fch = nn.Linear(hidden_size, output_size)  # 定义全连接层，用于输出
        self.fc = nn.Linear(output_size, output_size)  # 定义全连接层，用于输出


    def forward(self, x):
        # x: 输入数据，形状 (batch_size, seq_len, input_size)
        rnn_out, h_n = self.rnn(x)  # rnn_out: 所有时间步的输出，h_n: 最终时间步的隐藏状态

        out = self.fc(self.fch(rnn_out)) 

        out = torch.clamp(torch.relu(out),min=0,max=9)

        return out




